# encoding: utf-8
"""
@author:  shaoqijun
@date:  2020/7/26
"""
import os
import shutil
from config import cfg
import multiprocessing
import cv2
import time
import random
import numpy as np
import math
import numpy.random as npr
import logging
from utils.log_helper import init_log,add_file_handler
import glob

logger = logging.getLogger('global')


class gen_train_data():
    def __init__(self,img_list_txt, img_root, img_out_dir, bg_root):
        self.img_out_dir = img_out_dir
        self.img_list_txt = img_list_txt
        self.img_root = img_root
        self.bg_root = bg_root

        self._makedir(img_out_dir)
        self._check_before_run()
        self._gen_train_DataByMp()

    #运行前检查文件径路是否正确
    def _check_before_run(self):
        if not os.path.exists(self.img_list_txt):
            raise RuntimeError("'{}' is not available".format(self.img_list_txt))
        if not os.path.exists(self.img_root):
            raise RuntimeError("'{}' is not available".format(self.img_root))
        if not os.path.exists(self.bg_root):
            raise RuntimeError("'{}' is not available".format(self.bg_root))
        if not os.path.exists(self.img_out_dir):
            raise RuntimeError("'{}' is not available".format(self.img_out_dir))


    #创建文件夹
    def _makedir(self,dir):
        if not os.path.exists(dir):
            os.makedirs(dir)

    #获取背景图径路和数量
    def _get_bg(self):
        bg_files = []
        for i in os.listdir(cfg.DATA.BG_PATH):
            bg_files.append(os.path.join(cfg.DATA.BG_PATH,i))
        num_bg_files = len(bg_files)
        return bg_files, num_bg_files

    #获取图片名字列表和数量
    def _open_imglist_txt(self):
        img_sets = []
        with open(self.img_list_txt,'r') as f:
            for img_name in f.readlines():
                img_sets.append((img_name.strip()) )
        num_img_sets = len(img_sets)
        return img_sets, num_img_sets

    #数据增强
    def _train_data_aug(self, img_sets, start, end, bg_files, out_path):
        for i in range(start, end):
            print('{:.2f} %'.format((i - start) / (end - start) * 100))
            #print((i - start) / (end - start) * 100 , '%')
            img_name = img_sets[i]
            fold = img_name.strip().split('.')[0]
            self._makedir(os.path.join(out_path, fold))
            img_path = os.path.join(self.img_root, img_name)
            img = cv2.imread(img_path)
            for j in range(cfg.DATA.NUM_FOR_ROTATE):
                pi = random.randint(cfg.DATA.PI[0], cfg.DATA.PI[1])
                angle = random.randint(cfg.DATA.ANGLE[0], cfg.DATA.ANGLE[1])
                #time_start = time.time()
                aff_img = self._affine(img, pi, angle)
                #cv2.imshow("img",aff_img)
                #cv2.waitKey(0)
                #time_end = time.time()
                #print('time 1:', time_end - time_start, 's')

                time_start = time.time()

                bgImgs = []
                # for k in range(cfg.DATA.NUM_FOR_BG_CHOICE):
                #     num = random.randint(0, len(bg_files) - 1)
                #     bg_img = cv2.imread(bg_files[num])
                #     if bg_img is None:
                #         raise RuntimeError("file does not exits")
                #     bgImgs.append(bg_img)
                for bg_image in bg_files:
                    bg_img = cv2.imread(bg_image)
                    if bg_img is None:
                        raise RuntimeError("file does not exits")
                    bgImgs.append(bg_img)

                self._ori_shift_SF(aff_img, aff_img, bgImgs, out_path, fold, fold, pi, ori=False)

                time_end = time.time()
                #print('time 2:', time_end - time_start, 's')

    # 仿射变化
    def _affine(self, img, pi, angle):
        h, w, c = img.shape
        dist = 1  # 原相机到观测平面距离
        # pi = random.randint(0,30)
        pitch = pi * 3.14 / 180.0
        # angle = random.randint(0,360)
        roll = angle * 3.14 / 180.0
        R_AB = np.array([[math.cos(roll), math.sin(roll), 0], [-math.sin(roll), math.cos(roll), 0], [0, 0, 1]]).T
        R_BC = np.array([[1, 0, 0], [0, math.cos(pitch), math.sin(pitch)], [0, -math.sin(pitch), math.cos(pitch)]]).T
        R_AD = np.dot(R_AB, R_BC)
        t_AD = np.array([-math.sin(roll), math.cos(roll), 0]).reshape(3, 1)
        offset = dist * math.tan(pitch)  # 虚拟相机偏离原相机中心距离
        t_AD = t_AD * offset
        K = np.array([[w / 2, 0, w / 2], [0, w / 2, h / 2], [0, 0, 1]])
        R_DA = np.linalg.inv(R_AD)
        t_DA = -np.dot(R_DA, t_AD)
        n = np.array([0, 0, 1]).reshape(1, 3)
        nd = n / dist
        K_inv = np.linalg.inv(K)
        dot1 = np.dot(t_DA, nd)
        H = np.dot(K, np.dot((R_DA + np.dot(t_DA, nd)), K_inv))
        dst_img = cv2.warpPerspective(img, H, (int(w * 1), int(h * 1)),borderValue=(255,255,255))
        # pts = np.array([[0, 0, 1], [w - 1, 0, 1], [0, h - 1, 1], [w - 1, h - 1, 1]]).T
        # pts1 = np.dot(H1, pts)
        # pts1 = pts1 / pts1[2,:]
        # print(pts1)
        # cv2.imshow('image', img)
        # cv2.imshow('dst_image', dst_img)
        # cv2.waitKey()
        return dst_img

    # 变暗，加高斯噪声
    def _addNoise(self, img):
        intensity = random.randint(1, 5)
        noise = npr.randn(img.shape[0], img.shape[1], 3) * intensity
        scale = npr.uniform(0.4, 1.3)
        img_ = np.zeros(img.shape, dtype='uint8')
        height, width, n = img.shape
        for i in range(height):
            for j in range(width):
                for k in range(3):
                    c = int(img[i, j, k] * scale + noise[i, j, k])
                    if c > 255:
                        c = 255
                    if c < 0:
                        c = 0
                    img_[i, j, k] = c
        return img_

    #改变背景
    def _change_bg(self, bgImg, mask, font_ori):
        # cv2.imshow('mask', mask)
        # cv2.imshow('font_org', font_ori)
        # cv2.waitKey()
        x = npr.randint(0, bgImg.shape[1] - cfg.DATA.IMG_OUT_WIDTH)
        y = npr.randint(0, bgImg.shape[0] - cfg.DATA.IMG_OUT_HEIGHT)
        roi = bgImg[y: y + cfg.DATA.IMG_OUT_HEIGHT, x: x + cfg.DATA.IMG_OUT_WIDTH]
        # cv2.imshow('before', mask)
        mask = cv2.erode(mask, None, iterations=2)
        # cv2.imshow('erode', mask)
        mask = cv2.dilate(mask, None, iterations=1)
        # cv2.imshow('dilate', mask)

        mask3 = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
        new_img = np.where(mask3 == 0, roi, font_ori)
        # cv2.imshow('a.jpg',new_img)
        # cv2.waitKey()
        # new_img2gray = cv2.cvtColor(new_img, cv2.COLOR_BGR2GRAY)
        return new_img

    def _change_bg1(self,bgImg, mask, font_ori):
        bg_img_ = cv2.resize(bgImg, (326, 248), interpolation=cv2.INTER_CUBIC)
        # cv2.imshow('a.jpg',bg_img_)
        # cv2.waitKey()
        roi = bg_img_[0:160, 0:160]
        mask_inv = cv2.bitwise_not(mask)
        img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)
        img_fg = cv2.bitwise_and(font_ori, font_ori, mask=mask)
        dst = cv2.add(img_bg, img_fg)
        bg_img_[0:160, 0:160] = dst
        new_img = bg_img_[0:160, 0:160]
        erode = cv2.erode(new_img, None, iterations=1)
        # cv2.imshow('erode', erode)
        dilate = cv2.dilate(erode, None, iterations=1)
        # cv2.imshow('dilate', dilate)
        # cv2.imshow('a.jpg',img_fg)
        # cv2.waitKey()
        # new_img2gray = cv2.cvtColor(new_img, cv2.COLOR_BGR2GRAY)
        return dilate

    #旋转，偏移，背景变换，噪音
    def _ori_shift_SF(self, roi,src,bgImgs,outpath,fold,name,pi,ori=True):
        time_start = time.time()
        src_h = src.shape[0]
        src_w = src.shape[1]
        #cv2.imshow("roi",roi)
        gray_img = cv2.cvtColor(roi,cv2.COLOR_BGR2GRAY)
        ret, thresh = cv2.threshold(gray_img, 200, 255,  cv2.THRESH_BINARY_INV)
        # cv2.imshow("gr",gray_img)
        # cv2.imshow("th",thresh)
        # cv2.waitKey(0)
        gray_img = thresh
        contours, hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        w_h = []
        x_y = []
        d = []
        for i in range(len(contours)):
            rect = cv2.minAreaRect(contours[i])
            x1, y1 = rect[0]
            X_Y = (x1, y1)
            x_y.append(X_Y)
            width1, height1 = rect[1]
            W_H = (width1, height1)
            w_h.append(W_H)
            d.append(width1)
        index = d.index(max(d))
        x, y = x_y[index]
        width, height = w_h[index]

        time_end = time.time()
        # print('time 1:', time_end - time_start, 's')
        time_start = time.time()
        if ori == True:
            size = max(width, height) * 1.40
            x1_ = int(x - size / 2.)
            y1_ = int(y - size / 2.)
            x2_ = int(x1_ + size)
            y2_ = int(y1_ + size)
            print('width=', width, 'height=', height, 'x=', x, 'y=', y)
            if x1_ < 0 or y1_ < 0 or x2_ > src_w or y2_ > src_h:
                x1_1 = int(x - width / 2. - 1)
                y1_1 = int(y - height / 2. - 1)
                x2_1 = int(x + width / 2. + 1)
                y2_1 = int(y + height / 2. + 1)
                img_ = src[y1_1:y2_1, x1_1:x2_1:]
                gray_img_ = gray_img[y1_1:y2_1, x1_1:x2_1]
            else:
                img_ = src[y1_:y2_, x1_:x2_:]
                gray_img_ = gray_img[y1_:y2_, x1_:x2_]
            ori_img = cv2.resize(img_, (cfg.DATA.IMG_OUT_WIDTH, cfg.DATA.IMG_OUT_HEIGHT))
            new_gray_img = cv2.resize(gray_img_, (cfg.DATA.IMG_OUT_WIDTH, cfg.DATA.IMG_OUT_HEIGHT))
            n = random.randint(0, len(bgImgs) - 1)
            new_img = self._change_bg(bgImgs[n], new_gray_img, ori_img)
            # cv2.imshow('a.jpg',new_img)
            # cv2.waitKey()
            # cv2.imwrite(outpath + fold + '/{}ori.jpg'.format(name),ori_img)
            cv2.imwrite(outpath + fold + '/{}ori.jpg'.format(name), new_img)
            # cv2.imshow('a',ori_img)
            # cv2.waitKey()
        else:
            for i in range(cfg.DATA.NUM_FOR_TRANSLATION):
                size = npr.randint(int(max(width, height) * 1.1), np.ceil(1.6 * max(width, height)))
                delta_x = npr.randint(-width * 0.4, width * 0.4)
                delta_y = npr.randint(-height * 0.4, height * 0.4)
                left = int(x - width / 2. - 1)
                top = int(y - height / 2. - 1)
                right = int(x + width / 2. + 1)
                down = int(y + height / 2. + 1)
                x11 = int(max(int(x - size / 2. + delta_x / 2.), 0))
                y11 = int(max(int(y - size / 2. + delta_y / 2.), 0))
                x22 = int(x11 + size)
                y22 = int(y11 + size)
                if x22 > src_w or y22 > src_h or x22 < right or y22 < down or x11 > left or y11 > top:
                    continue
                img_ = src[y11:y22, x11:x22:]
                gray_img_ = gray_img[y11:y22, x11:x22]
                new_gray_img = cv2.resize(gray_img_, (cfg.DATA.IMG_OUT_WIDTH, cfg.DATA.IMG_OUT_HEIGHT))
                img_h = img_.shape[0]
                img_w = img_.shape[1]
                # print(img_h,img_w)
                ori_img = cv2.resize(img_, (cfg.DATA.IMG_OUT_WIDTH, cfg.DATA.IMG_OUT_HEIGHT))
                # cv2.imshow("ori_img", ori_img)

                time_start1 = time.time()
                n = random.randint(0, len(bgImgs) - 1)
                new_img = self._change_bg(bgImgs[n], new_gray_img, ori_img)
                # cv2.imshow("new_img", new_img)
                # cv2.waitKey()
                new_img = self._addNoise(new_img)
                time_end1 = time.time()
                #print('time change:', time_end1 - time_start1, 's')
                # cv2.imwrite(outpath + fold + '/{}w{}_h{}_{}.jpg'.format(name,img_w,img_h,num), ori_img)
                cv2.imwrite(os.path.join(outpath, fold, '{}w{}_h{}_{}_pi{}.png').format(name, img_w, img_h, i, pi), new_img)
                #cv2.imwrite(outpath + fold + '/{}w{}_h{}_{}_pi{}.png'.format(name, img_w, img_h, i, pi), new_img)
        time_end = time.time()
        #print('time 2:', time_end - time_start, 's')

    #cv2读图
    def _read_img(self,img):
        if os.path.exists(img):
            img = cv2.imread(img)
            cv2.imshow("img", img)
            cv2.waitKey(1)

    #多线程数据增强
    def _gen_train_DataByMp(self):
        img_sets, num_img_sets = self._open_imglist_txt()
        bg_files, num_bg_sets = self._get_bg()

        index = []
        pn = cfg.DATA.PN
        pool = multiprocessing.Pool(processes=pn)
        for i in range(pn):
            index.append(int(i*len(img_sets)/pn))
        index.append(len(img_sets))


        for i in range(pn):
            #continue
            pool.apply_async(self._train_data_aug,
                             (img_sets, index[i], index[i + 1], bg_files, self.img_out_dir))

        pool.close()
        pool.join()

        logger.info("=> button  loaded")
        logger.info("Dataset statistics:")
        logger.info("---------------------------------")
        logger.info("     file        |   !@#$%^&*   | # numbers")
        logger.info("---------------------------------")
        logger.info("num for img list |              |    {}".format(num_img_sets))
        logger.info("num for bg       |              |    {}".format(num_bg_sets))
        logger.info("num for pn       |              |    {}".format(pn))

        total_num = 0
        total_ids = 0
        for fold in img_sets:
            total_ids+=1
            fold = fold.strip().split('.')[0]
            num = self._get_num_for_fold(fold)
            total_num += int(num)
        logger.info("num for total    |\t\t|   {}".format(total_num))
        logger.info("num for ids      |\t\t|   {}".format(total_ids))
        logger.info("---------------------------------")
        print('finished')

    def _get_num_for_fold(self,fold):
        fold_path = os.path.join(cfg.DATA.OUT_PATH,fold)
        if os.path.exists(fold_path):
            file_name = os.path.join(fold_path, "*.png")
            path_file_number = glob.glob(pathname=file_name)
            num_in_fold = len(path_file_number)
            logger.info("num for {}  |\t\t|  {}".format(fold,num_in_fold))
            return num_in_fold
        return 0


def get_log():
    if os.path.exists(cfg.TRAIN.LOG_DIR):
        shutil.rmtree(cfg.TRAIN.LOG_DIR)
    os.makedirs(cfg.TRAIN.LOG_DIR)
    init_log('global', logging.INFO)
    if cfg.TRAIN.LOG_DIR:
        add_file_handler('global', os.path.join(cfg.TRAIN.LOG_DIR, "data_logs.txt"), logging.info)

def main():
    get_log()
    logger.info("for data aug")
    img_list_txt = cfg.DATA.IMG_LIST_ROOT
    img_root = cfg.IMGLIST.IMG_IN_DIR
    out_dir = cfg.DATA.OUT_PATH
    bg_root = cfg.DATA.BG_PATH
    gen_train_data(img_list_txt, img_root, out_dir, bg_root)


if __name__ == '__main__':
    main()